import os
import shutil
from pathlib import Path

def read_label_file(label_path):
    """读取标签文件，返回类别ID"""
    with open(label_path, 'r') as f:
        # YOLO格式的标签文件第一个数字是类别ID
        class_id = int(f.readline().split()[0])
    return class_id

def get_unique_classes(root_dir):
    """扫描所有标签文件，获取所有唯一的类别ID"""
    unique_classes = set()
    for split in ['train', 'test', 'valid']:
        labels_dir = os.path.join(root_dir, split, 'labels')
        if not os.path.exists(labels_dir):
            continue
            
        for label_file in os.listdir(labels_dir):
            if label_file.endswith('.txt'):
                label_path = os.path.join(labels_dir, label_file)
                try:
                    class_id = read_label_file(label_path)
                    unique_classes.add(class_id)
                except Exception as e:
                    print(f"Error reading {label_file}: {e}")
    return sorted(list(unique_classes))

def organize_dataset(root_dir, target_dir):
    # 自动获取所有类别
    unique_classes = get_unique_classes(root_dir)
    print(f"Found classes: {unique_classes}")
    
    # 创建类别映射
    class_names = {i: f'class{i}' for i in unique_classes}
    print(f"Class mapping: {class_names}")
    
    # 处理训练集、验证集和测试集
    for split in ['train', 'test', 'valid']:
        if not os.path.exists(os.path.join(root_dir, split)):
            continue
            
        images_dir = os.path.join(root_dir, split, 'images')
        labels_dir = os.path.join(root_dir, split, 'labels')
        
        # 为每个类别创建目标目录
        for class_name in class_names.values():
            os.makedirs(os.path.join(target_dir, split, class_name), exist_ok=True)
        
        # 处理每个图片文件
        for img_file in os.listdir(images_dir):
            img_path = os.path.join(images_dir, img_file)
            label_file = os.path.join(labels_dir, img_file.rsplit('.', 1)[0] + '.txt')
            
            if os.path.exists(label_file):
                try:
                    # 读取类别ID并获取类别名称
                    class_id = read_label_file(label_file)
                    class_name = class_names[class_id]
                    
                    # 复制图片到对应的类别目录
                    dst_path = os.path.join(target_dir, split, class_name, img_file)
                    shutil.copy2(img_path, dst_path)
                    print(f"Copied {img_file} to {class_name} directory")
                except Exception as e:
                    print(f"Error processing {img_file}: {e}")

if __name__ == '__main__':
    # 源数据集路径
    root_dir = "F:/Desktop/wjx_bishe/YOLO_format"
    # 目标数据集路径（建议使用新的目录）
    target_dir = "F:/Desktop/wjx_bishe/YOLO_format_cls"
    
    organize_dataset(root_dir, target_dir) 